{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 处理代码加载\n",
    "import torch\n",
    "from transformers import AutoProcessor, LlavaForConditionalGeneration,LlavaProcessor\n",
    "\n",
    "model = LlavaForConditionalGeneration.from_pretrained('/data/VLM/llava-v1.5-7b-hf', torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, attn_implementation=\"eager\").to('cuda')\n",
    "processor = LlavaProcessor.from_pretrained('/data/VLM/llava-v1.5-7b-hf', patch_size=14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from loader import *\n",
    "image_path = '/data/dataset/M3FD/M3FD_Detection/Vis/00000.png'\n",
    "lr_image_path = '/data/dataset/M3FD/M3FD_Detection/Ir/00000.png'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageDraw\n",
    "from run import vicrop_qa\n",
    "\n",
    "model_name = 'llava'\n",
    "method_name = 'rel_att'\n",
    "\n",
    "question = 'Identify the positions of the cars in the picture.'\n",
    "short_question = 'Identify the positions of the cars in the picture.'\n",
    "\n",
    "# Run the Vicrop method\n",
    "ori_answer, crop_answer, bbox = vicrop_qa(model_name, method_name, image_path, question, model, processor, short_question)\n",
    "lr_ori_answer, lr_crop_answer, lr_bbox = vicrop_qa(model_name, method_name, lr_image_path, question, model, processor, short_question)\n",
    "print(f'Model\\'s original answer:  {ori_answer}')\n",
    "print(f'Answer with Vicrop:       {crop_answer}')\n",
    "\n",
    "# Visualize the bounding box\n",
    "image = Image.open(image_path).convert(\"RGB\")\n",
    "lr_image = Image.open(lr_image_path).convert(\"RGB\")\n",
    "image_draw = ImageDraw.Draw(image)\n",
    "image_draw.rectangle(bbox, outline='red', width=4)\n",
    "image_draw.rectangle(lr_bbox, outline='blue', width=4)\n",
    "display(image.resize((500, 500*image.size[1]//image.size[0])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llava",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
